import json

from langchain.chat_models import init_chat_model
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool


@tool(description="计算2个int类型的和")
def add(a: int, b: int) -> int:
    return a + b


@tool(description="计算2个float类型的乘积")
def multiply(a: float, b: float) -> float:
    return a * b


llm = init_chat_model("deepseek-chat", model_provider="deepseek",
                      )

llm_with_tools = llm.bind_tools([add, multiply])

query = "3.5的4倍是多少"
# query = "1加2是多少"
messages = [HumanMessage(query)]
output = llm_with_tools.invoke(messages)
messages.append(output)
tool_call_func = {"multiply": multiply, "add": add}
for tool_call in output.tool_calls:
    select_func = tool_call_func.get(tool_call["name"])
    tool_msg = select_func.invoke(tool_call)
    messages.append(tool_msg)

new_output = llm_with_tools.invoke(messages)
print(new_output.content)

for message in messages:
    print(message)
